"""
Yuxi-Know MCP工具定义
将Yuxi-Know核心功能封装为MCP工具供外部系统调用
"""
from typing import Dict, List, Optional, Any, Union
from pydantic import BaseModel, Field
from pathlib import Path
import json
import httpx


class KnowledgeQueryInput(BaseModel):
    """知识库查询输入参数"""
    database_id: str = Field(description="知识库ID", example="kb_123456")
    query: str = Field(description="查询内容", example="人工智能技术的发展历史")
    top_k: int = Field(default=5, description="返回结果数量", ge=1, le=20)
    search_type: str = Field(default="hybrid", description="检索类型：local/global/hybrid/naive/mix",
                           pattern="^(local|global|hybrid|naive|mix)$")
    threshold: float = Field(default=0.7, description="相似度阈值", ge=0.0, le=1.0)


class GraphQueryInput(BaseModel):
    """图谱查询输入参数"""
    node_label: str = Field(description="图谱节点标签", example="人工智能")
    depth: int = Field(default=2, description="查询深度", ge=1, le=5)
    max_nodes: int = Field(default=100, description="最大节点数量", ge=1, le=1000)


class ChatInput(BaseModel):
    """对话输入参数"""
    agent_id: str = Field(description="智能体ID", example="chatbot")
    message: str = Field(description="用户消息", example="请介绍一下机器学习的基本概念")
    thread_id: Optional[str] = Field(default=None, description="会话线程ID（可选）")
    knowledge_base_ids: Optional[List[str]] = Field(default=None, description="关联的知识库ID列表")
    model_name: Optional[str] = Field(default=None, description="指定模型名称")


class KnowledgeBaseCreateInput(BaseModel):
    """知识库创建输入参数"""
    name: str = Field(description="知识库名称", example="人工智能知识库")
    description: str = Field(description="知识库描述", example="包含AI领域相关文档和资料")
    kb_type: str = Field(default="LightRAG", description="知识库类型：LightRAG/Chroma/Milvus",
                        pattern="^(LightRAG|Chroma|Milvus)$")
    embedding_model: str = Field(default="bge-m3", description="嵌入模型名称")
    llm_info: Optional[Dict[str, Any]] = Field(default=None, description="LLM配置信息")


class SystemHealthOutput(BaseModel):
    """系统健康状态输出"""
    status: str = Field(description="系统状态")
    message: str = Field(description="状态描述")
    timestamp: str = Field(description="检查时间")


class KnowledgeQueryOutput(BaseModel):
    """知识库查询结果输出"""
    results: List[Dict[str, Any]] = Field(description="查询结果列表")
    count: int = Field(description="结果数量")
    time_taken: float = Field(description="查询耗时（秒）")
    database_name: str = Field(description="知识库名称")


class GraphDataOutput(BaseModel):
    """图谱数据输出"""
    nodes: List[Dict[str, Any]] = Field(description="节点列表")
    edges: List[Dict[str, Any]] = Field(description="边列表")
    total_nodes: int = Field(description="总节点数")
    total_edges: int = Field(description="总边数")
    depth_exceeded: bool = Field(description="是否超出最大深度限制")


class ChatOutput(BaseModel):
    """对话输出格式"""
    response: str = Field(description="智能体回复")
    thread_id: str = Field(description="会话线程ID")
    agent_id: str = Field(description="使用的智能体ID")
    timestamp: str = Field(description="响应时间")
    context_used: bool = Field(description="是否使用了知识上下文")


class AvailableKnowledgeBasesOutput(BaseModel):
    """可用知识库列表"""
    databases: List[Dict[str, Any]] = Field(description="知识库列表")
    total_count: int = Field(description="知识库总数")
    active_count: int = Field(description="活跃知识库数量")


class EntitySearchInput(BaseModel):
    """实体搜索输入"""
    entity_name: str = Field(description="实体名称", example="深度学习")
    limit: int = Field(default=10, description="返回结果数量", ge=1, le=50)


class EntitySearchOutput(BaseModel):
    """实体搜索结果"""
    entities: List[Dict[str, Any]] = Field(description="实体列表")
    relationships: List[Dict[str, Any]] = Field(description="关系列表")
    total_count: int = Field(description="实体总数")


class EntityRelationsInput(BaseModel):
    """实体关系查询输入"""
    entity_name: str = Field(description="实体名称", example="贾敏")
    limit: int = Field(default=20, description="返回关系数量", ge=1, le=50)


class EntityRelationsOutput(BaseModel):
    """实体关系查询结果"""
    entity_name: str = Field(description="查询的实体名称")
    relations_summary: str = Field(description="关系汇总（自然语言描述）")
    detailed_relations: List[str] = Field(description="详细关系列表（自然语言描述）")
    total_relations: int = Field(description="关系总数")


class DocumentUploadInput(BaseModel):
    """文档上传输入"""
    database_id: str = Field(description="目标知识库ID")
    documents: List[str] = Field(description="文档文件路径列表")
    urls: Optional[List[str]] = Field(default=None, description="URL链接列表")


class DocumentUploadOutput(BaseModel):
    """文档上传结果"""
    task_id: str = Field(description="处理任务ID")
    status: str = Field(description="处理状态")
    message: str = Field(description="状态描述")
    documents_processed: int = Field(description="已处理文档数量")


class ToolListOutput(BaseModel):
    """可用工具列表输出"""
    tools: List[Dict[str, Any]] = Field(description="工具列表")
    categories: List[str] = Field(description="工具分类")
    total_count: int = Field(description="工具总数")


class GraphStatsOutput(BaseModel):
    """图谱统计信息输出"""
    total_nodes: int = Field(description="总节点数")
    total_edges: int = Field(description="总边数")
    entity_types: Dict[str, int] = Field(description="实体类型分布")
    last_updated: str = Field(description="最后更新时间")


# MCP工具函数定义
async def get_system_health() -> SystemHealthOutput:
    """
    获取Yuxi-Know系统健康状态

    返回值包含系统运行状态、消息和时间戳
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.check_system_health()

        return result
    except Exception as e:
        # 如果API调用失败，返回错误状态
        from datetime import datetime
        return SystemHealthOutput(
            status="error",
            message=f"健康检查失败: {str(e)}",
            timestamp=datetime.now().isoformat()
        )


async def query_knowledge_base(input_data: KnowledgeQueryInput) -> KnowledgeQueryOutput:
    """
    查询指定知识库中的内容

    支持多种检索模式：local（局部）、global（全局）、hybrid（混合）、naive（朴素）、mix（混合）
    可根据相似度阈值和时间进行内容筛选
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.search_knowledge_base(input_data)

        return result
    except Exception as e:
        # 如果API调用失败，返回空结果
        return KnowledgeQueryOutput(
            results=[],
            count=0,
            time_taken=0.0,
            database_name=input_data.database_id,
        )


async def query_knowledge_graph(input_data: GraphQueryInput) -> GraphDataOutput:
    """
    查询知识图谱，获取相关子图结构

    可以指定查询深度和最大节点数量，返回相关的知识图谱网络
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.query_graph(input_data)

        return result
    except Exception as e:
        # 如果API调用失败，返回空结果
        return GraphDataOutput(
            nodes=[],
            edges=[],
            total_nodes=0,
            total_edges=0,
            depth_exceeded=False
        )


async def chat_with_agent(input_data: ChatInput) -> ChatOutput:
    """
    与AI智能体进行对话

    支持多种智能体类型，可以关联知识库进行上下文增强
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.chat_with_intelligent_agent(input_data)

        return result
    except Exception as e:
        # 如果API调用失败，返回错误消息
        from datetime import datetime
        return ChatOutput(
            response=f"对话服务异常: {str(e)}",
            thread_id=input_data.thread_id or "error_thread",
            agent_id=input_data.agent_id,
            timestamp=datetime.now().isoformat(),
            context_used=False
        )


async def create_knowledge_base(input_data: KnowledgeBaseCreateInput) -> Dict[str, Any]:
    """
    创建新的知识库

    支持多种知识库类型和配置选项
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.create_knowledge_base_internal(input_data)

        return result
    except Exception as e:
        # 如果API调用失败，返回错误信息
        return {
            "error": f"创建知识库失败: {str(e)}",
            "status_code": 500,
            "success": False
        }


async def get_knowledge_bases() -> AvailableKnowledgeBasesOutput:
    """
    获取所有可用的知识库列表

    返回知识库的详细信息和元数据
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.get_available_knowledge_bases()

        return result
    except Exception as e:
        # 如果API调用失败，返回空结果
        return AvailableKnowledgeBasesOutput(
            databases=[],
            total_count=0,
            active_count=0
        )


async def search_entities(input_data: EntitySearchInput) -> EntitySearchOutput:
    """
    在知识图谱中搜索实体

    支持模糊匹配和向量相似度搜索
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.search_graph_entities(input_data)

        return result
    except Exception as e:
        # 如果API调用失败，返回空结果
        return EntitySearchOutput(
            entities=[],
            relationships=[],
            total_count=0
        )


async def get_available_tools() -> ToolListOutput:
    """
    获取系统中可用的工具列表

    包括知识检索、计算、搜索等各种工具
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.get_system_tools()

        return result
    except Exception as e:
        # 如果API调用失败，返回空结果
        return ToolListOutput(
            tools=[],
            categories=[],
            total_count=0
        )


async def get_graph_statistics() -> GraphStatsOutput:
    """
    获取知识图谱的统计信息

    包括节点数量、类型分布等
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.get_graph_statistics_info()

        return result
    except Exception as e:
        # 如果API调用失败，返回空结果
        from datetime import datetime
        return GraphStatsOutput(
            total_nodes=0,
            total_edges=0,
            entity_types={},
            last_updated=datetime.now().isoformat()
        )


async def query_entity_relations(input_data: EntityRelationsInput) -> EntityRelationsOutput:
    """
    查询指定实体的所有关系，以自然语言描述返回

    输入实体名称，返回该实体的所有关系的自然语言描述
    """
    from .mcp_server import get_mcp_server

    try:
        # 获取MCP服务器实例并调用实际API
        mcp_server = get_mcp_server()
        result = await mcp_server.get_entity_relations(input_data)

        return result
    except Exception as e:
        # 如果API调用失败，返回空结果
        return EntityRelationsOutput(
            entity_name=input_data.entity_name,
            relations_summary=f"查询实体'{input_data.entity_name}'的关系时发生错误: {str(e)}",
            detailed_relations=[],
            total_relations=0
        )


# 工具定义映射用于注册
TOOLS_MAP = {
    "yuxi-know-health": {
        "function": get_system_health,
        "name": "yuxi-know-health",
        "description": "获取Yuxi-Know知识系统健康状态",
        "input_schema": None,  # 无输入参数
        "output_schema": SystemHealthOutput
    },
    "yuxi-know-query-knowledge": {
        "function": query_knowledge_base,
        "name": "yuxi-know-query-knowledge",
        "description": "在Yuxi-Know知识库中查询相关内容，支持多种检索模式",
        "input_schema": KnowledgeQueryInput,
        "output_schema": KnowledgeQueryOutput
    },
    "yuxi-know-query-graph": {
        "function": query_knowledge_graph,
        "name": "yuxi-know-query-graph",
        "description": "查询知识图谱，获取相关子图结构和关联关系",
        "input_schema": GraphQueryInput,
        "output_schema": GraphDataOutput
    },
    "yuxi-know-chat": {
        "function": chat_with_agent,
        "name": "yuxi-know-chat",
        "description": "与AI智能体进行对话，支持上下文关联和知识库增强",
        "input_schema": ChatInput,
        "output_schema": ChatOutput
    },
    "yuxi-know-create-kb": {
        "function": create_knowledge_base,
        "name": "yuxi-know-create-kb",
        "description": "创建新的知识库，支持多种类型和配置选项",
        "input_schema": KnowledgeBaseCreateInput,
        "output_schema": Dict[str, Any]
    },
    "yuxi-know-get-kbs": {
        "function": get_knowledge_bases,
        "name": "yuxi-know-get-kbs",
        "description": "获取所有可用的知识库列表和详细信息",
        "input_schema": None,  # 无输入参数
        "output_schema": AvailableKnowledgeBasesOutput
    },
    "yuxi-know-search-entities": {
        "function": search_entities,
        "name": "yuxi-know-search-entities",
        "description": "在知识图谱中搜索实体和关联关系",
        "input_schema": EntitySearchInput,
        "output_schema": EntitySearchOutput
    },
    "yuxi-know-get-tools": {
        "function": get_available_tools,
        "name": "yuxi-know-get-tools",
        "description": "获取系统中可用的工具列表",
        "input_schema": None,  # 无输入参数
        "output_schema": ToolListOutput
    },
    "yuxi-know-graph-stats": {
        "function": get_graph_statistics,
        "name": "yuxi-know-graph-stats",
        "description": "获取知识图谱的统计分析信息",
        "input_schema": None,  # 无输入参数
        "output_schema": GraphStatsOutput
    },
    "yuxi-know-entity-relations": {
        "function": query_entity_relations,
        "name": "yuxi-know-entity-relations",
        "description": "查询指定实体的所有关系，返回自然语言描述",
        "input_schema": EntityRelationsInput,
        "output_schema": EntityRelationsOutput
    }
}